
#include "group_norm_v2_tiling.h"
#include "register/op_def_registry.h"

using namespace std;

#define ull uint64_t

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    // 获取属性的数量和值
    int attr_num = context->GetAttrs()->GetAttrNum();
    const int64_t *num_groups0 = context->GetAttrs()->GetInt(0);
    ull num_groups = *num_groups0;
    bool data_format = 0;  // NCHW
    float eps = 0.0001;
    bool is_training = true;
    if (attr_num > 1) {
        const char *data_format0 = context->GetAttrs()->GetStr(1);
        if (strcmp(data_format0, "NCHW") != 0) {
            data_format = 1;  // NHWC
        }
        const float *eps0 = context->GetAttrs()->GetFloat(2);
        eps = *eps0;
        const bool *is_training0 = context->GetAttrs()->GetBool(3);
        is_training = *is_training0;
    }

    // 获取输入tensor的形状
    auto shape_x = context->GetInputShape(0)->GetOriginShape();
    // 有多少个维度
    int dimNum = shape_x.GetDimNum();
    // 每个维度的大小
    ull n = shape_x.GetDim(0);
    ull c = shape_x.GetDim(1);
    ull hw = 1;
    for (int i = 2; i < dimNum; i++) {
        hw *= shape_x.GetDim(i);
    }
    ull totalLength = n * c * hw;

    // 获取输入tensor的数据类型
    auto tensor = context->GetInputTensor(0);
    auto type = tensor->GetDataType();
    if(type==ge::DT_FLOAT16) {
        context->SetTilingKey(1);
        context->SetBlockDim(1);
    } else {
        context->SetTilingKey(2);
        context->SetBlockDim(n * num_groups);
    }

    GroupNormV2TilingData tiling;
    tiling.set_totalLength(totalLength);
    tiling.set_tileNum(1);
    tiling.set_num_groups(num_groups);
    tiling.set_data_format(data_format);
    tiling.set_eps(eps);
    tiling.set_is_training(is_training);
    tiling.set_n(n);
    tiling.set_c(c);
    tiling.set_hw(hw);

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class GroupNormV2 : public OpDef {
public:
    explicit GroupNormV2(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("gamma")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("beta")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("mean")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("rstd")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("num_groups").Int();
        this->Attr("data_format").AttrType(OPTIONAL).String("NCHW");
        this->Attr("eps").AttrType(OPTIONAL).Float(0.0001);
        this->Attr("is_training").AttrType(OPTIONAL).Bool(true);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(GroupNormV2);
}
